/** \file    refine/bit_mask.cpp
  * \brief   Moves across the mask and adds or removes single pixels to it 
  *			 according to color ditribution in image.
  * \date    2003 - 2005
  * \author  Fedor Chelnokov
  * \version $Header: /home/cscom1r4/cvsroot/cs/refine/bit_mask.cpp,v 1.16 2006/01/15 07:31:32 cscom1r4 Exp $
  */

#include "refine/sorted_neighbours.h"
#include "refine/bit_mask.h"
#include "imglib/imageclass.h"
#include "imglib/bitmap.h"
#include "imglib/bmpmorph.h"
#include <iostream>
#include <limits>
#include <list>
#include <math.h>
using std::clog;
using std::endl;

namespace refine {

namespace bit_mask {

params::params()
{
	whs = 6;
	auto_t = true;
	t = 3.5f;
	max_passes = 150;
	consider = info_all;
	change = turn_on_and_off;
	prefer = prefer_none;
	silent = false;

	modifiable = 0;
}

bool operator == (const params & lhs, const params & rhs)
{
	if (lhs.whs != rhs.whs ||
		lhs.auto_t != rhs.auto_t ||
		lhs.max_passes != rhs.max_passes ||
		lhs.consider != rhs.consider ||
		lhs.change != rhs.change ||
		//lhs.silent != rhs.silent ||
		lhs.modifiable != rhs.modifiable)
		return false;

	if (!lhs.auto_t) {
		if (lhs.t != rhs.t ||
			lhs.prefer != rhs.prefer)
			return false;
	}

	return true;
}

////////////////////
// Core functions //
////////////////////

enum classification {
	class_0,
	class_1,
	class_denial
};

inline void complex_op(bool b, bool & and, bool & or) {
	and = and && b;
	or = or || b;
}

// prior classification is done using mask values only:
// if all neighbours belong to one class then there is no other choice
classification prior(const mask & m, size_point const & c) {
	size_t w = m.width();
	size_t h = m.height();
	bool mul = true, sum = false;
	if (c.x > 0 && c.y > 0)
		complex_op(m(c.x-1, c.y-1), mul, sum);
	if (c.x > 0)
		complex_op(m(c.x-1, c.y  ), mul, sum);
	if (c.x > 0 && c.y + 1 < h)
		complex_op(m(c.x-1, c.y+1), mul, sum);
	if (c.y + 1 < h)
		complex_op(m(c.x  , c.y+1), mul, sum);
	if (c.x + 1 < w && c.y + 1 < h)
		complex_op(m(c.x+1, c.y+1), mul, sum);
	if (c.x + 1 < w)
		complex_op(m(c.x+1, c.y  ), mul, sum);
	if (c.x + 1 < w && c.y > 0)
		complex_op(m(c.x+1, c.y-1), mul, sum);
	if (c.y > 0)
		complex_op(m(c.x  , c.y-1), mul, sum);
	if (mul)
		return class_1;
	if (!sum)
		return class_0;
	return class_denial;
}

template <class R>
classification auto_classify(sorted_neighbours<R> & sort_cache, const mask & m, 
	size_point const & c, icolor const & col)
{
	uint nei_size = std::numeric_limits<uint>::max();
	const neighbours & nei = sort_cache.get(c, col, nei_size);
	assert (nei.size() == nei_size);

	classification res = class_denial;
	float best_t = 0; //so far

	int d = 0;
	for (uint n = 0; n != nei.size(); ++n) {
		const bool nei_state = m(c.x + nei[n].second.x, c.y + nei[n].second.y);
		if (nei_state)
			++d;
		else
			--d;

		//plateau (several equal color distances) process as single point
		if (n+1 != nei.size() && nei[n].first == nei[n+1].first)
			continue;

		if (d * d > best_t * (n+1)) {
			best_t = float(d * d) / (n+1);
			res = (d > 0) ? class_1 : class_0;
		}
	}
	return res;
}

template <class R>
classification classify(sorted_neighbours<R> & sort_cache, const mask & m, 
	size_point const & c, icolor const & col, const params & p)
{
	if (p.auto_t)
		return auto_classify(sort_cache, m, c, col);

	assert (p.t > 1);
	uint nei_sorted_size = 2 * p.t;
	const neighbours & nei = sort_cache.get(c, col, nei_sorted_size);

	int d = 0;
	for (uint n = 0; n != nei.size(); ++n) {
		if (n+1 >= nei_sorted_size) {
			nei_sorted_size *= 2;
			sort_cache.get(c, col, nei_sorted_size);
		}
		const bool nei_state = m(c.x + nei[n].second.x, c.y + nei[n].second.y);
		if (nei_state)
			++d;
		else
			--d;

		//plateau (several equal color distances) process as single point
		if (n+1 != nei.size() && nei[n].first == nei[n+1].first)
			continue;

		if (d * d > p.t * (n+1))
			return (d > 0) ? class_1 : class_0;
	}
	return class_denial;
}

template <class R>
void core(
	const R & ref_im, const mask & ref_m, 
	const R & tgt_im, mask & tgt_m, 
	const affine * motion,
	const params & p, bool check_once
)
{
	size_extent ref_dim = ref_im.dim();
	assert(ref_dim == ref_m.dim());
	size_extent tgt_dim = tgt_im.dim();
	assert(tgt_dim == tgt_m.dim());
	if (ref_m.data() != tgt_m.data())
		check_once = true; //since classification of any pixel depends only from ref_m, which does not change
	sorted_neighbours<R> sort_cache(ref_im, p.whs, p.consider);

	if (!p.silent)
		clog << "refining initiated..." << endl;

	mask modifiable(tgt_dim);
	modifiable.fill(true);

	mask not_checked_yet;
	if (check_once) {
		not_checked_yet.resize(tgt_dim);
		not_checked_yet.fill(true);
	}

	mask next_m = tgt_m;

	affine back_proj;
	if (motion) {
		back_proj = motion->inverse();
		back_proj.translate(0.5f, 0.5f); //for proper rounding later
	} else
		back_proj.identity();
	int_rect ref_frame(0, 0, ref_dim.sx, ref_dim.sy);

	for (int pass = 0; ; ++pass) 
	{
		if (check_once)
			img::combine(modifiable, not_checked_yet, img::bit_and());
		if (p.modifiable)
			img::combine(modifiable, *p.modifiable, img::bit_and());

		uint checks = 0, changes = 0, scans = sort_cache.scans(), sorts = sort_cache.sorts();

		for (mask::iterator c = modifiable.begin(); c != modifiable.end(); ++c) {
			if (!*c)
				continue;
			bool tgt_b = tgt_m[c];
			if (p.change == turn_on_only && tgt_b || p.change == turn_off_only && !tgt_b) {
				*c = false;
				continue;
			}

			classification cl = prior(next_m, c);
			if (cl == class_denial) {
				++checks;
				if (check_once)
					not_checked_yet[c] = false;
				//back project c
				int_point ref_c = back_proj((size_point)c);
				if (ref_frame.inside(ref_c))
					cl = classify(sort_cache, ref_m, ref_c, tgt_im[c], p);
				if (p.prefer != prefer_none && cl == class_denial) 
					cl = (p.prefer == prefer_masked_pixels) ? class_1 : class_0;
			}
			if (cl == class_0 && !tgt_b || cl == class_1 && tgt_b)
				cl = class_denial;
			if (cl == class_denial)
				*c = false;
			else {
				next_m[c] = !tgt_b;
				++changes;
			}
		}

		scans = sort_cache.scans() - scans;
		sorts = sort_cache.sorts() - sorts;
		if (!p.silent)
			clog << "pass " << pass <<
				": checks=" << checks << 
				", cache_size=" << sort_cache.size() <<
				", sorts=" << sorts << 
				", scans=" << scans << 
				", changes=" << changes << endl;

		if (changes == 0 || scans == 0 && sorts == 0 //sorts&scans are zero if some cycling takes place
			|| pass >= p.max_passes)
			break;

		tgt_m = next_m;

		// modified on the current pass points can influence state of other points
		// on the distance up to p.whs
		if (p.whs <= 2)
			img::dilate8(modifiable, p.whs); //for square window
		else
			img::dilate4(modifiable, p.whs); //for circle window (slightly less than needed)
		sort_cache.pass_done();
	}
}

////////////////////////
// Exported functions //
////////////////////////

void adjust(
	const byte_image & im, mask & m, 
	const params & p, bool check_once)
{
	core(im, m, im, m, 0, p, check_once);
}

void adjust(
	const argb_byte_image & im, mask & m, 
	const params & p, bool check_once)
{
	core(im, m, im, m, 0, p, check_once);
}

void track(
	const byte_image & ref_im, const mask & ref_m, 
	const byte_image & tgt_im, mask & tgt_m, 
	information_source isource,
	const affine * motion,
	const params & p
)
{
	if (isource == source_reference || isource == source_both) {
		if (!p.silent)
			clog << "\nrefine using reference frame\n";
		core(ref_im, ref_m, tgt_im, tgt_m, motion, p, true);
	}
	if (isource == source_target || isource == source_both) {
		if (!p.silent)
			clog << "\nrefine (adjust)\n";
		adjust(tgt_im, tgt_m, p, false);
	}
}

} //namespace bit_mask

} //namespace refine
